import numpy as np
import networkx as nx
from collections import defaultdict
from typing import Dict, List, Any, Tuple
from app.schemas.dhgam.mplp_schema import MPLPInputParams, MPLPOutputParams, AlgorithmResponse, AlgorithmMiddleResponse, TaskStatus, InputParams, OutputParams
from app.utils.logger import logger


# 元路径链接预测服务 (MPLP)
class MetaPathLinkPrediction:
    def __init__(self, meta_paths=None):
        self.meta_paths = meta_paths or []
        self.path_counts = defaultdict(float)
    
    def add_meta_path(self, path):
        """添加元路径定义, 如 ['User', 'Item', 'User']"""
        self.meta_paths.append(path)
    
    def compute_meta_path_counts(self, graph, source_type, target_type):
        """计算元路径计数"""
        self.path_counts = defaultdict(float)
        
        # 获取所有源节点和目标节点
        source_nodes = [n for n in graph.nodes() if graph.nodes[n]['type'] == source_type]
        target_nodes = [n for n in graph.nodes() if graph.nodes[n]['type'] == target_type]
        
        # 为每个元路径计算路径计数
        for path in self.meta_paths:
            for source in source_nodes:
                # 使用BFS沿着元路径搜索
                current_nodes = [source]
                for i, node_type in enumerate(path[1:], 1):
                    next_nodes = []
                    for current in current_nodes:
                        for neighbor in graph.neighbors(current):
                            if graph.nodes[neighbor]['type'] == node_type:
                                next_nodes.append(neighbor)
                    current_nodes = next_nodes
                
                # 计算最终到达的目标节点
                for target in current_nodes:
                    if target in target_nodes:
                        self.path_counts[(source, target)] += 1
    
    def predict(self, graph, source_type, target_type, top_k=10):
        """预测可能的链接"""
        self.compute_meta_path_counts(graph, source_type, target_type)
        
        # 获取所有源节点和目标节点
        source_nodes = [n for n in graph.nodes() if graph.nodes[n]['type'] == source_type]
        target_nodes = [n for n in graph.nodes() if graph.nodes[n]['type'] == target_type]
        
        # 计算链接预测分数
        predictions = []
        for source in source_nodes:
            for target in target_nodes:
                # 跳过已存在的边
                if graph.has_edge(source, target):
                    continue
                    
                score = self.path_counts.get((source, target), 0)
                if score > 0:
                    predictions.append((source, target, score))
        
        # 返回分数最高的 top_k 个预测
        predictions.sort(key=lambda x: x[2], reverse=True)
        return predictions[:top_k]


class MPLPService:
    def __init__(self):
        pass
        
    def _build_graph_from_input(self, nodes: List[Dict[str, Any]], edges: List[Dict[str, Any]]) -> nx.Graph:
        """
        从输入参数构建NetworkX图
        
        Args:
            nodes: 节点列表 [{"id": id, "type": type, ...}, ...]
            edges: 边列表 [{"source": source_id, "target": target_id, ...}, ...]
            
        Returns:
            nx.Graph: 构建的NetworkX图
        """
        G = nx.Graph()
        
        # 添加节点
        for node in nodes:
            node_id = node.pop("id")
            G.add_node(node_id, **node)
            
        # 添加边
        for edge in edges:
            source = edge.pop("source")
            target = edge.pop("target")
            G.add_edge(source, target, **edge)
            
        return G
        
    async def run_mplp(self, input_params: Dict[str, Any], task_id: str) -> AlgorithmResponse:
        """
        运行元路径链接预测算法
        
        Args:
            input_params: 算法输入参数
            task_id: 任务ID
            
        Returns:
            AlgorithmResponse: 算法执行结果
        """
        try:
            # 解析输入参数
            mplp_params = input_params.get("mplp_params")
            if not mplp_params:
                raise ValueError("缺少MPLP参数")
                
            # 构建图
            nodes = mplp_params.get("nodes", [])
            edges = mplp_params.get("edges", [])
            meta_paths = mplp_params.get("meta_paths", [])
            source_type = mplp_params.get("source_type")
            target_type = mplp_params.get("target_type")
            top_k = mplp_params.get("top_k", 10)
            
            # 验证必要参数
            if not nodes or not edges or not meta_paths:
                raise ValueError("缺少节点、边或元路径定义")
            if not source_type or not target_type:
                raise ValueError("缺少源节点类型或目标节点类型")
                
            # 构建图
            graph = self._build_graph_from_input(nodes, edges)
            
            # 创建MPLP实例并运行算法
            mplp = MetaPathLinkPrediction()
            for path in meta_paths:
                mplp.add_meta_path(path)
                
            predictions = mplp.predict(graph, source_type, target_type, top_k)
            
            # 格式化预测结果
            formatted_predictions = [
                {"source": source, "target": target, "score": float(score)}
                for source, target, score in predictions
            ]
            
            # 格式化路径计数
            path_counts_dict = {f"{s}_{t}": float(v) for (s, t), v in mplp.path_counts.items()}
            
            # 构建输出
            mplp_output = {
                "predictions": formatted_predictions,
                "algorithm": "MetaPathLinkPrediction",
                "parameters": {
                    "meta_paths": meta_paths,
                    "source_type": source_type,
                    "target_type": target_type,
                    "top_k": top_k
                },
                "path_counts": path_counts_dict
            }
            
            # 构建响应
            response = AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.COMPLETED,
                task_progress=100,
                output_params=OutputParams(
                    mplp_results=MPLPOutputParams(**mplp_output)
                ),
                metrics=[]  # 可以添加可视化图表
            )
            
            return response
            
        except Exception as e:
            logger.error(f"MPLP执行错误: {str(e)}")
            return AlgorithmResponse(
                task_id=task_id,
                task_status=TaskStatus.FAILED,
                error_message=str(e),
                output_params=OutputParams(),
                metrics=[]
            ) 